import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.feature_selection import RFE
from statsmodels.stats.outliers_influence import variance_inflation_factor
from common_import import *


# 计算VIF值的函数
def calculate_vif(df):
    vif_data = pd.DataFrame()
    vif_data["Feature"] = df.columns
    vif_data["VIF"] = [
        (
            variance_inflation_factor(df.values, i)
            if df.iloc[:, i].nunique() > 1
            else float("inf")
        )
        for i in range(df.shape[1])
    ]
    return vif_data


# 结合VIF的RFE特征选择函数
def rfe_with_vif(data, vif_threshold=10):
    X = data
    selected_features = X.columns.tolist()
    # print("Initial Features:", selected_features)
    vif_values = []
    while True:
        # 计算当前选中的特征的 VIF

        X_selected = X[selected_features]
        if X_selected.isnull().values.any() or np.isinf(X_selected.values).any():
            X_selected = X_selected.replace([np.inf, -np.inf], np.nan).dropna()
        vif_data = calculate_vif(X_selected)

        # 处理VIF值为无穷大的情况，将其设为一个非常大的值
        # vif_data["VIF"].replace([np.inf, -np.inf], 1e10, inplace=True)

        # 输出当前最大VIF值
        max_vif_value = max(vif_data["VIF"])
        vif_values.append(min(max_vif_value, 10000))
        print(f"Current max VIF: {max_vif_value}")

        # 检查所有特征的VIF值是否都小于阈值
        if max_vif_value < vif_threshold:
            break

        # 找到VIF最高的特征
        max_vif_feature = vif_data.sort_values("VIF", ascending=False).iloc[0][
            "Feature"
        ]
        print(f"Removing feature: {max_vif_feature} with VIF: {max_vif_value}")

        # 移除该特征
        selected_features.remove(max_vif_feature)

    return selected_features, vif_values


if __name__ == "__main__":
    # 读取包含前关注特征的文件
    focus_feature = constants.feature_295[:135]
    # 读取包含所有特征的数据文件
    data = pd.read_csv("data/Molecular_Descriptor_training.csv")
    data_focus = data[focus_feature]
    # data_focus_standardized = data_focus
    print(focus_feature[:63])
    # 标准化数据
    # data_focus_standardized = (data_focus - data_focus.mean()) / data_focus.std()
    # 假设target为一个y值，可以用你实际的目标变量替换这里
    # target = data_focus["target_column"]  # 请替换为你的目标列
    # 如果没有具体目标变量，只进行特征选择
    # result_features, vif_values = rfe_with_vif(data_focus_standardized)
    # plt.figure(figsize=(10, 6))
    # plt.plot(vif_values, marker="o")
    # plt.xlabel("迭代次数")
    # plt.ylabel("当前最大VIF值")
    # plt.title("")
    # plt.xlim(left=0)
    # plt.ylim([0, 10100])  # 设置y轴范围，将超过1000的部分截断
    # plt.grid(True)
    # # plt.show()
    # tool.show_or_print("vif.png")
    print(result_features)
    print(len(result_features))
    # 显示筛选后的特征
    # import ace_tools as tools

    # tools.display_dataframe_to_user(
    #     name="Filtered Features",
    #     dataframe=pd.DataFrame(result_features, columns=["Selected Features"]),
    # )
